{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.137298583984375, -0.49993896484375, -3.41259765625, 0.364166259765625, -0.29769134521484375, 1.70263671875, 0.37176513671875, -1.1875, -0.59844970703125, -0.88525390625, 0.24542236328125, 1.208984375, 0.190277099609375, 1.9697265625, 0.82568359375, -1.05224609375, 0.6501235961914062, -0.612548828125, -0.9190673828125, 0.73223876953125, 0.149261474609375, -1.632080078125, -0.54248046875, 0.8173828125, 1.9300537109375, -0.04766082763671875, -0.301055908203125, 1.5107421875, 0.1049652099609375, -0.160064697265625, -0.1656036376953125, -0.248809814453125, 0.22821044921875, 0.71917724609375, 0.82666015625, 0.370849609375, 0.81683349609375, 0.3165550231933594, 0.7071533203125, 0.0792999267578125, 0.15337657928466797, -0.939697265625, 0.0965576171875, 0.36815643310546875, 1.1243896484375, 0.290863037109375, 0.29449462890625, 0.133087158203125, 1.835693359375, -0.8916015625, -0.5359115600585938, -0.124267578125, -0.2336273193359375, 1.389892578125, 2.03955078125, 0.5783615112304688, 0.945556640625, -1.06634521484375, 0.26197052001953125, 0.9561767578125, 0.06524658203125, 0.9017333984375, 0.891845703125, 1.5614013671875, -0.0746917724609375, -0.87158203125, -0.524200439453125, 0.34716796875, -0.5392379760742188, 0.85662841796875, 1.717529296875, 0.53045654296875, 0.0620880126953125, 0.2812459468841553, 0.9136962890625, 0.7674560546875, 0.8658447265625, 0.423004150390625, -0.0344390869140625, 0.051544189453125, 1.212158203125, -0.26837158203125, 1.083984375, -1.12109375, -0.4675140380859375, -1.519287109375, -1.0064697265625, 0.6748046875, -1.823486328125, -0.508514404296875, 1.7025146484375, 0.56121826171875, 0.3792877197265625, -0.278900146484375, -0.9677734375, 0.6114501953125, -1.290771484375, -0.2029876708984375, -0.397674560546875, 0.036846160888671875, 0.1953125, -0.828369140625, 0.22882080078125, 1.533203125, 1.982421875, 0.88134765625, -1.195556640625, -0.9156494140625, -1.0587158203125, -0.8382568359375, -0.8406982421875, 0.001560211181640625, 0.14934587478637695, 0.0291748046875, -0.13306045532226562, -0.7435302734375, 0.4718208312988281, -1.74169921875, 0.4305000305175781, 0.494232177734375, -0.29736328125, 0.02919769287109375, -0.763519287109375, 0.04746246337890625, 0.78466796875, 0.64306640625, -1.693359375, -0.3915863037109375, 0.594482421875, 0.0334625244140625, 0.59283447265625, -0.4635009765625, -0.903564453125, -0.5738525390625, -1.03955078125, 0.37725830078125, -0.771240234375, 0.304443359375, 1.009033203125, 0.77081298828125, 0.04810905456542969, 0.12504005432128906, 0.653167724609375, -1.29443359375, -1.108642578125, -0.884521484375, 1.4461669921875, -1.549560546875, 0.52374267578125, -0.17523193359375, 0.464324951171875, 0.9959716796875, 0.7603759765625, 1.282958984375, 0.5450439453125, -0.396575927734375, -0.151123046875, 0.388580322265625, 1.0811767578125, -1.2596435546875, 0.3662567138671875, -0.042205810546875, 0.055469512939453125, 0.14417266845703125, 0.64593505859375, -1.696044921875, 0.3870849609375, 1.431396484375, 0.05583381652832031, 0.3656730651855469, -0.90838623046875, -0.8260498046875, 0.66143798828125, -0.90869140625, -0.76904296875, -0.227783203125, -0.385772705078125, -0.45832061767578125, 0.5511474609375, -0.1149444580078125, 0.2137451171875, -0.7447509765625, 0.309234619140625, 0.15543079376220703, -0.0319976806640625, -0.61480712890625, -0.76080322265625, 0.71392822265625, 0.22430419921875, -0.738037109375, -0.4757232666015625, 0.9239501953125, -1.2437744140625, -0.7362060546875, 0.6256103515625, -0.4614410400390625, 0.3103914260864258, 0.2319488525390625, -0.2533607482910156, -1.9256591796875, -1.14166259765625, 0.3807373046875, -1.1339111328125, -1.0430908203125, -0.1565093994140625, 1.71044921875, 1.414794921875, 1.334228515625, 0.533935546875, 0.13201904296875, 1.060546875, -1.27294921875, -1.31103515625, -0.57666015625, -0.2191009521484375, 0.034759521484375, -0.13318252563476562, -1.1644287109375, 0.118560791015625, -0.7613525390625, 0.34930419921875, -0.35477447509765625, 1.403564453125, 0.47247314453125, 1.2996826171875, 0.1222381591796875, -1.5478515625, -0.07476806640625, -1.0946044921875, 0.67047119140625, -0.74737548828125, -1.588134765625, 0.033843994140625, 0.51220703125, -0.259185791015625, 0.983642578125, 0.7276611328125, 1.062255859375, 0.19046783447265625, -0.58428955078125, -0.6337890625, 0.5170745849609375, 0.405517578125, -0.867919921875, -0.6741943359375, 0.18458938598632812, 0.9901123046875, -0.1063232421875, -0.202423095703125, 0.2523040771484375, -0.9625244140625, -0.6253662109375, -0.339080810546875, -0.1998291015625, 0.7835693359375, 0.17961883544921875, -0.2379150390625, -0.2716522216796875, 0.65478515625, -0.2747802734375, 0.310516357421875, -2.8984375, 1.4033203125, -0.521392822265625, -0.70245361328125, 0.875244140625, -0.78466796875, 0.84625244140625, 0.782470703125, -0.875732421875, 0.451446533203125, 0.26373291015625, -0.22841835021972656, 0.58636474609375, -0.774139404296875, 0.0893707275390625, 0.3702392578125, 0.3561859130859375, -0.0710601806640625, 0.69818115234375, -0.076141357421875, -0.55352783203125, 0.033341407775878906, -0.782470703125, 0.3951416015625, -0.066864013671875, -0.2907257080078125, 0.34003448486328125, -0.62237548828125, -0.10138702392578125, 0.334991455078125, 0.79443359375, 0.7515869140625, 1.9150390625, 0.256103515625, 1.036376953125, 0.5159912109375, 0.4652099609375, 0.44842529296875, 0.2004241943359375, -0.2091217041015625, 0.166961669921875, 1.248779296875, 0.2104339599609375, -0.4805908203125, 0.24781036376953125, 1.771484375, -0.8040771484375, 0.593994140625, 0.10785675048828125, -0.746826171875, 1.95947265625, -2.36767578125, -0.25003814697265625, -1.626708984375, 0.806640625, -0.1574249267578125, 0.526092529296875, 1.425048828125, 0.590576171875, -0.8355712890625, -1.2181396484375, 0.6962890625, -0.315185546875, -0.2504425048828125, 0.6372604370117188, 0.459716796875, -0.04080963134765625, 0.396575927734375, -1.3448486328125, -0.90509033203125, 0.333648681640625, -0.736572265625, -0.37823486328125, -0.10791015625, 0.757720947265625, 0.28856658935546875, -0.78662109375, -1.1328125, 0.9949951171875, -0.3224639892578125, 0.302276611328125, 0.8206787109375, -0.383209228515625, -0.413116455078125, -0.27384185791015625, -0.52288818359375, 0.21628570556640625, -0.001861572265625, 0.78173828125, -0.370849609375, 0.521209716796875, -1.01416015625, 0.35655975341796875, -0.10810089111328125, 0.446868896484375, 0.1050262451171875, 0.55694580078125, 0.0038003921508789062, -1.337890625, -0.4964599609375, 0.18148422241210938, 0.1795201301574707, -0.66192626953125, -0.3319129943847656, 0.42232513427734375, 0.892333984375, 0.93212890625, 0.51513671875, -0.7236328125, -0.33127593994140625, -0.51953125, 0.241119384765625, 0.3442535400390625, 0.22637939453125, -0.260894775390625, -0.880126953125, -0.298065185546875, -0.00970458984375, 0.9273681640625, 0.010848045349121094, -0.08880615234375, 0.52362060546875, -0.3914031982421875, -0.01934814453125, 0.6119384765625, 0.24212646484375, -0.6859130859375, 0.33514404296875, 0.77276611328125, 0.33441162109375, -0.333953857421875, 0.205291748046875, -0.2896156311035156, -0.33526611328125, -0.4931640625, -0.589874267578125, -0.658935546875, -0.5668411254882812, 0.557861328125, -0.20977783203125, -0.2099609375, 0.47808837890625, -1.39794921875, 0.1097564697265625, -0.539306640625, -1.218017578125, -0.61376953125, 0.28521728515625, -0.062164306640625, 1.557861328125, 1.5859375, 0.903564453125, -1.555908203125, -0.606201171875, -0.12288665771484375, 0.7364501953125, -1.44384765625, -0.481414794921875, 0.06729316711425781, -0.07932186126708984, 1.156494140625, 0.001678466796875, 0.13311767578125, -0.46319580078125, 1.16650390625, 0.29547119140625, 0.425323486328125, 0.067779541015625, -0.615234375, -0.46307373046875, 0.5543804168701172, 0.0718994140625, 0.264190673828125, -0.4052734375, -0.46148681640625, 1.2607421875, -0.26763916015625, 0.9595947265625, 0.86737060546875, 1.523193359375, -0.91943359375, -1.29248046875, -0.06358623504638672, 0.267425537109375, 2.659912109375, 1.47314453125, 0.06403541564941406, -1.5732421875, -0.062450408935546875, -0.68896484375, -0.13484954833984375, 1.287109375, -0.240570068359375, 2.160400390625, -1.1102294921875, 0.5892333984375, -0.23095703125, -0.0772705078125, 1.218017578125, -0.704833984375, 0.267303466796875, -1.164794921875, 0.3043060302734375, 0.6744384765625, -0.617523193359375, -0.231719970703125, 0.122833251953125, 0.10870361328125, 0.14825439453125, -0.478271484375, -0.822998046875, -0.49371337890625, -0.7130126953125, -1.32470703125, 0.0604095458984375, -0.6739501953125, -0.4495201110839844, -0.7534866333007812, 0.607666015625, 1.034912109375, 0.17634201049804688, -1.501953125, 0.084136962890625, 1.074462890625, 2.28173828125, -0.203216552734375, 0.20513916015625, -0.67315673828125, 0.008575439453125, 0.016594409942626953, -0.13214111328125, 0.7884521484375, 0.2082958221435547, -0.8341064453125, 0.5590362548828125, -0.48016357421875, 1.0280685424804688, -0.33795166015625, -0.18111419677734375, 1.4219970703125, -0.469024658203125, 0.057476043701171875, 1.0880126953125, 0.37139892578125, 0.6461181640625, 0.48309326171875, -1.3515625, -0.3076629638671875, 0.17486572265625, 1.1187744140625, -0.32952880859375, -0.487884521484375, 0.4190673828125, 2.0654296875, -0.85260009765625, 0.54010009765625, -0.5982666015625, -0.37200927734375, 0.7174100875854492, 0.02947998046875, 0.284637451171875, -1.0579833984375, 0.12664031982421875, -1.633544921875, 1.2998046875, -0.29360198974609375, -1.0985107421875, 0.5392570495605469, 0.0411834716796875, 0.349029541015625, 0.3655548095703125, -0.05511474609375, -1.266357421875, 1.0341796875, -0.56646728515625, 0.33941650390625, 0.060028076171875, 0.562103271484375, 0.256866455078125, -0.390594482421875, -0.66693115234375, 0.11818885803222656, -1.1279296875, 1.2623291015625, 0.95208740234375, -1.559814453125, 0.315460205078125, 1.180908203125, 0.4473114013671875, -2.028564453125, -0.7822265625, -0.7496337890625, 0.06356430053710938, -0.9442138671875, -0.69732666015625, 0.0055389404296875, 0.5057945251464844, -0.670135498046875, -1.65283203125, 0.823486328125, 1.684326171875, -0.90936279296875, 0.2619171142578125, 0.5441360473632812, 0.011810302734375, -0.1874103546142578, -0.1158294677734375, 0.12662506103515625, -1.1322021484375, -0.0003662109375, -0.8778076171875, 0.611053466796875, -0.79290771484375, 0.26214599609375, 0.2493896484375, 1.1512451171875, -0.466552734375, -0.65924072265625, -0.378631591796875, -0.2699737548828125, -1.23974609375, 0.10634613037109375, 0.547760009765625, -0.1894378662109375, 0.48974609375, -0.0997772216796875, -1.45751953125, 0.713165283203125, -0.375152587890625, -0.5706787109375, 0.75146484375, 0.8284912109375, -0.74066162109375, -0.498779296875, 1.383056640625, -0.6795806884765625, -1.469329833984375, -0.239776611328125, -0.83416748046875, -0.50640869140625, 0.052093505859375, 1.78857421875, -1.969482421875, 0.84967041015625, 1.78466796875, -0.34747314453125, 0.9990234375, 0.017425537109375, 0.13677978515625, 1.137451171875, 0.621826171875, -0.5199737548828125, 0.0911865234375, 0.67919921875, -0.21998214721679688, 1.3101806640625, -0.0619354248046875, -0.6424560546875, -2.07470703125, -1.1788330078125, 0.3248138427734375, 0.86602783203125, 0.522216796875, 0.18072509765625, -0.05080413818359375, -1.211181640625, -0.98291015625, -1.090576171875, 1.071533203125, 0.6453857421875, -0.2294464111328125, -1.584716796875, 0.464324951171875, 0.22609710693359375, 0.288726806640625, 0.13329124450683594, 0.368621826171875, -1.445068359375, 0.38675689697265625, 1.3756103515625, 0.4884033203125, 0.4590606689453125, -0.1070556640625, 0.5357666015625, -1.379150390625, 0.982421875, 0.8924560546875, 1.0069580078125, 0.8568115234375, 1.546630859375, 1.5947265625, 1.0732421875, -0.7099609375, -0.1714324951171875, -0.42572021484375, 0.015228271484375, -0.697998046875, -2.292236328125, 0.62286376953125, 0.056488037109375, 0.027778148651123047, -0.8995361328125, -0.66070556640625, 0.7723388671875, -0.921875, -0.4859619140625, 0.4466552734375, -1.946533203125, 0.54559326171875, -0.5518798828125, -0.70751953125, -0.661376953125, -1.236083984375, -0.4008979797363281, 0.9698486328125, 0.88702392578125, 0.878173828125, 0.095184326171875, 0.2110137939453125, 0.5194034576416016, -0.35650634765625, -0.78460693359375, 0.588134765625, 0.6986083984375, -0.449005126953125, 1.32080078125, -0.9534912109375, -1.5, -1.255859375, -0.60693359375, -0.25286865234375, 0.13799285888671875, -1.164794921875, -0.6710205078125, 0.44482421875, 0.03527069091796875, 0.6502685546875, -1.1412353515625, 0.738037109375, -0.827880859375, -0.1855621337890625, 0.48614501953125, -0.0853271484375, -0.0193328857421875, -0.986572265625, 0.29712772369384766, 0.5576171875, 0.098297119140625, -0.010284423828125, -0.350616455078125, 0.7808837890625, 0.259490966796875, 1.414306640625, -0.2137451171875, 1.369384765625, -0.741668701171875, -0.7626953125, -0.137298583984375, 1.0623779296875, 1.6265869140625, 1.47900390625, -0.3381805419921875, -0.66229248046875, -0.76806640625, -0.19713592529296875, 0.4137420654296875, -0.7615966796875, 0.5374755859375, -0.5428466796875, -0.7381591796875, -0.7352294921875, -0.81451416015625, 0.3705902099609375, -0.985107421875, 0.7325439453125, -0.221893310546875, -1.172119140625, 0.2322540283203125, -0.54498291015625, -0.8543701171875, -0.85308837890625, 0.49053955078125, 1.0540771484375, 0.10843658447265625, -0.37982177734375, 0.34454345703125, 2.31689453125, 0.276031494140625, -0.05600547790527344, 0.3253326416015625, -0.198577880859375, 0.05181884765625, -1.838134765625, -0.2581787109375, 0.410980224609375, 0.5791015625, 0.2361602783203125, 1.27978515625, -0.7127685546875, 0.55389404296875, -0.45111083984375, 0.08251953125, -0.16156005859375, 0.917236328125, -0.76312255859375, -1.1876220703125, -0.312957763671875]\n"
     ]
    }
   ],
   "source": [
    "from langchain_community.chat_models import ChatOllama\n",
    "from langchain_core.messages import AIMessage, HumanMessage,SystemMessage\n",
    "llm2str = ChatOllama(model=\"qwen2:7b\", temperature=0)\n",
    "from langchain_ollama import OllamaEmbeddings\n",
    "\n",
    "embedding = OllamaEmbeddings(model=\"nomic-embed-text:v1.5\")\n",
    "print(embedding.embed_query(\"Hello world\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|im_start|>system\n",
      "你的任务是在多轮对话中，生成一个改写后的用户问题，以补全上下文中缺失的信息，一般的技巧是指代消岐和信息补全。<|im_end|>\n",
      "<|im_start|>user\n",
      "今天上海的天气怎么样\n",
      "<|im_end|>\n",
      "<|im_start|>assistant\n",
      "今天上海 18 度，晴天\n",
      "<|im_end|>\n",
      "<|im_start|>user\n",
      "明天呢？\n",
      "<|im_end|>\n",
      "<|im_start|>assistant\n",
      "明天预计 19 度，阴天\n",
      "<|im_end|>\n",
      "-----请改写下面的用户问题，以补全上下文中缺失的信息，一般的技巧是指代消岐和信息补全-----\n",
      "<|im_start|>user\n",
      "广州什么天气？<|im_end|>\n",
      "原问题:广州什么天气？\n",
      "改写后的问题:广州今天的天气怎么样？\n",
      "0.977381115909202\n",
      "{'rewritten': False, 'rewrite_query': '广州今天的天气怎么样？'}\n"
     ]
    }
   ],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "\n",
    "from flask import Flask\n",
    "from flask import request\n",
    "from flask import jsonify\n",
    "# coding=UTF-8\n",
    " \n",
    "\n",
    "\n",
    "# res = client.chat.completions.create(model=qwen14b_modelId, **request.get_json())\n",
    "\n",
    "\n",
    "\n",
    "# 在这里实现你的策略\n",
    "def digest_mrc(messages):\n",
    "    \"\"\"\n",
    "    # Request:\n",
    "    {\n",
    "    \"messages\": [\n",
    "        {\"role\": \"user\", \"content\": \"今天上海的天气怎么样\"},\n",
    "        {\"role\": \"system\", \"content\": \"今天上海 18 度，晴天\"},\n",
    "        {\"role\": \"user\", \"content\": \"明天呢？\"},\n",
    "        {\"role\": \"system\", \"content\": \"明天预计 19 度，阴天\"},\n",
    "        {\"role\": \"user\", \"content\": \"广州什么天气？\"}\n",
    "    ]\n",
    "    }\n",
    "\n",
    "    # Response:\n",
    "    {\n",
    "    \"rewritten\": true/false\n",
    "    \"rewrite_query\": \"广州明天什么天气？\"\n",
    "    }\n",
    "    # 如果需要改写就是true，然后query就是改写后的问题，否则就是false，query字段可有可无\n",
    "    \"\"\"\n",
    "    rewritten = False\n",
    "    rewrite_query = messages[-1][\"content\"]\n",
    "    # 策略 1 直接语言模型 提示改写\n",
    "    ## 先将message处理成合适的prompt\n",
    "    prompt_list = [\"<|im_start|>system\\n你的任务是在多轮对话中，生成一个改写后的用户问题，以补全上下文中缺失的信息，一般的技巧是指代消岐和信息补全。<|im_end|>\"]\n",
    "    query = messages[-1][\"content\"]\n",
    "    messages.pop(-1)\n",
    "    for message in messages:\n",
    "        if message[\"role\"] == \"user\":\n",
    "            prompt_list.append(\"<|im_start|>user\\n\" + message[\"content\"] + \"\\n<|im_end|>\")\n",
    "        elif message[\"role\"] == \"system\":\n",
    "            prompt_list.append(\"<|im_start|>assistant\\n\" + message[\"content\"] + \"\\n<|im_end|>\")\n",
    "        else:\n",
    "            raise ValueError(\"Invalid role: \" + message[\"role\"])\n",
    "    prompt_list.append(\"-----请改写下面的用户问题，以补全上下文中缺失的信息，一般的技巧是指代消岐和信息补全-----\\n<|im_start|>user\\n\" + query + \"<|im_end|>\")\n",
    "    prompt = \"\\n\".join(prompt_list)\n",
    "    print(prompt)\n",
    "    # print(prompt)\n",
    "    # res = client.chat.completions.create(model=qwen14b_modelId, prompt=prompt, temperature=0.0, max_tokens=100)\n",
    "    # if res.choices[0].message.content.strip() != query:\n",
    "    #     rewritten = True\n",
    "    #     rewrite_query = res.choices[0].message.content.strip()\n",
    "    rewrite_query = llm2str.invoke([{\"role\": \"user\", \"content\": prompt}]).content\n",
    "    rewrite_query_emb = embedding.embed_query(rewrite_query)\n",
    "    query_emb = embedding.embed_query(query)\n",
    "    from sklearn.metrics.pairwise import cosine_similarity\n",
    "    import numpy as np\n",
    "    print(\"原问题:\"+query)\n",
    "    print(\"改写后的问题:\"+rewrite_query)\n",
    "    embeddings = np.array([rewrite_query_emb, query_emb])\n",
    "    similarity_matrix = cosine_similarity(embeddings)\n",
    "    print(similarity_matrix[0, 1])\n",
    "    return {\n",
    "            \"rewritten\": rewritten,\n",
    "            \"rewrite_query\": rewrite_query\n",
    "        }\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    messages=[\n",
    "        {\"role\": \"user\", \"content\": \"今天上海的天气怎么样\"},\n",
    "        {\"role\": \"system\", \"content\": \"今天上海 18 度，晴天\"},\n",
    "        {\"role\": \"user\", \"content\": \"明天呢？\"},\n",
    "        {\"role\": \"system\", \"content\": \"明天预计 19 度，阴天\"},\n",
    "        {\"role\": \"user\", \"content\": \"广州什么天气？\"}\n",
    "    ]\n",
    "    print(digest_mrc(messages))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "# 第一次策略，没有判断是否需要重写。\n",
    "from flask import Flask\n",
    "from flask import request\n",
    "from flask import jsonify\n",
    "# coding=UTF-8\n",
    "import json\n",
    "import os\n",
    "from flask import make_response\n",
    "from openai import OpenAI\n",
    " \n",
    " \n",
    "model_info = json.loads(os.getenv(\"LEADERBOARD_MODELHUB_KEY2INFO\")) # 环境变量名称固定\n",
    "token = model_info[\"token\"] # 访问modelhub模型的token，固定key\n",
    "base_url = model_info[\"entrypoint\"] # 访问modelhub模型的base url，固定key\n",
    "qwen72b_modelId = model_info[\"model_key2info\"][\"qwen72b\"][\"modelId\"] # key的qwen72b与提交配置里的自定义名称对应\n",
    "client = OpenAI(api_key=token, base_url=base_url)\n",
    "\n",
    "# res = client.chat.completions.create(model=qwen72b_modelId, **request.get_json())\n",
    "\n",
    "app = Flask(__name__)\n",
    "\n",
    "\n",
    "# 在这里实现你的策略\n",
    "def digest_mrc(messages):\n",
    "    \"\"\"\n",
    "    # Request:\n",
    "    {\n",
    "    \"messages\": [\n",
    "        {\"role\": \"user\", \"content\": \"今天上海的天气怎么样\"},\n",
    "        {\"role\": \"system\", \"content\": \"今天上海 18 度，晴天\"},\n",
    "        {\"role\": \"user\", \"content\": \"明天呢？\"},\n",
    "        {\"role\": \"system\", \"content\": \"明天预计 19 度，阴天\"},\n",
    "        {\"role\": \"user\", \"content\": \"广州什么天气？\"}\n",
    "    ]\n",
    "    }\n",
    "\n",
    "    # Response:\n",
    "    {\n",
    "    \"rewritten\": true/false\n",
    "    \"rewrite_query\": \"广州明天什么天气？\"\n",
    "    }\n",
    "    # 如果需要改写就是true，然后query就是改写后的问题，否则就是false，query字段可有可无\n",
    "    \"\"\"\n",
    "    rewritten = False\n",
    "    rewrite_query = messages[-1][\"content\"]\n",
    "    # 策略 1 直接语言模型 提示改写\n",
    "    ## 先将message处理成合适的prompt\n",
    "    prompt_list = [\"<|im_start|>system\\n你的任务是根据多轮查询问题，将最新的问题改写成让用户只关注该问题就能理解这个问题<|im_end|>\"]\n",
    "    query = messages[-1][\"content\"]\n",
    "    messages.pop(-1)\n",
    "    for message in messages:\n",
    "        if message[\"role\"] == \"user\":\n",
    "            prompt_list.append(\"<|im_start|>user\\n\" + message[\"content\"] + \"\\n<|im_end|>\")\n",
    "        elif message[\"role\"] == \"system\":\n",
    "            prompt_list.append(\"<|im_start|>assistant\\n\" + message[\"content\"] + \"\\n<|im_end|>\")\n",
    "        else:\n",
    "            raise ValueError(\"Invalid role: \" + message[\"role\"])\n",
    "    prompt_list.append(\"<|im_start|>user\\n-----请改写下面这个问题，让用户只关注该问题就能理解这个问题-----\\n\" + query + \"<|im_end|>\")\n",
    "    prompt = \"\\n\".join(prompt_list)\n",
    "    \n",
    "    try:\n",
    "        completion = client.chat.completions.create(model=qwen72b_modelId, messages=[{\"role\": \"user\", \"content\": prompt}])\n",
    "\n",
    "        rewritten_query = completion.choices[0].message.content\n",
    "        if rewritten_query != query:\n",
    "            rewritten = True\n",
    "            rewrite_query = rewritten_query\n",
    "    except Exception as e:\n",
    "        print(f\"Error during OpenAI API call: {e}\")\n",
    "        return {\n",
    "            \"rewritten\": rewritten,\n",
    "            \"rewrite_query\": query\n",
    "        }\n",
    "    return {\n",
    "        \"rewritten\": rewritten,\n",
    "        \"rewrite_query\": rewrite_query\n",
    "    }\n",
    "\n",
    "\n",
    "# 此处不需要修改\n",
    "@app.route(\"/predict\",methods=[\"POST\"])\n",
    "def predict():\n",
    "    data = request.get_json()\n",
    "    messages = data.get('messages')\n",
    "\n",
    "    return digest_mrc(messages)\n",
    "\n",
    "\n",
    "# 此处不需要修改\n",
    "if __name__ == '__main__':\n",
    "   app.run(\"0.0.0.0\",80)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "# 第二次策略，判断是否需要重写查询。向量相似度判断\n",
    "from flask import Flask\n",
    "from flask import request\n",
    "from flask import jsonify\n",
    "# coding=UTF-8\n",
    "import json\n",
    "import os\n",
    "from flask import make_response\n",
    "from openai import OpenAI\n",
    " \n",
    " \n",
    "model_info = json.loads(os.getenv(\"LEADERBOARD_MODELHUB_KEY2INFO\")) # 环境变量名称固定\n",
    "token = model_info[\"token\"] # 访问modelhub模型的token，固定key\n",
    "base_url = model_info[\"entrypoint\"] # 访问modelhub模型的base url，固定key\n",
    "qwen72b_modelId = model_info[\"model_key2info\"][\"qwen72b\"][\"modelId\"] # key的qwen72b与提交配置里的自定义名称对应\n",
    "client = OpenAI(api_key=token, base_url=base_url)\n",
    "embedding_model = model_info[\"model_key2info\"][\"embedding\"][\"modelId\"]\n",
    "# res = client.chat.completions.create(model=qwen72b_modelId, **request.get_json())\n",
    "\n",
    "app = Flask(__name__)\n",
    "\n",
    "\n",
    "# 在这里实现你的策略\n",
    "def digest_mrc(messages):\n",
    "    \"\"\"\n",
    "    # Request:\n",
    "    {\n",
    "    \"messages\": [\n",
    "        {\"role\": \"user\", \"content\": \"今天上海的天气怎么样\"},\n",
    "        {\"role\": \"system\", \"content\": \"今天上海 18 度，晴天\"},\n",
    "        {\"role\": \"user\", \"content\": \"明天呢？\"},\n",
    "        {\"role\": \"system\", \"content\": \"明天预计 19 度，阴天\"},\n",
    "        {\"role\": \"user\", \"content\": \"广州什么天气？\"}\n",
    "    ]\n",
    "    }\n",
    "\n",
    "    # Response:\n",
    "    {\n",
    "    \"rewritten\": true/false\n",
    "    \"rewrite_query\": \"广州明天什么天气？\"\n",
    "    }\n",
    "    # 如果需要改写就是true，然后query就是改写后的问题，否则就是false，query字段可有可无\n",
    "    \"\"\"\n",
    "    rewritten = False\n",
    "    rewrite_query = messages[-1][\"content\"]\n",
    "    # 策略 1 直接语言模型 提示改写\n",
    "    ## 先将message处理成合适的prompt\n",
    "    prompt_list = [\"<|im_start|>system\\n你的任务是根据多轮查询问题，将最新的问题改写成让用户只关注该问题就能理解这个问题<|im_end|>\"]\n",
    "    query = messages[-1][\"content\"]\n",
    "    context = \"\"\n",
    "    messages.pop(-1)\n",
    "    for message in messages:\n",
    "        context += message[\"content\"] + \"\\n\"\n",
    "        if message[\"role\"] == \"user\":\n",
    "            prompt_list.append(\"<|im_start|>user\\n\" + message[\"content\"] + \"\\n<|im_end|>\")\n",
    "        elif message[\"role\"] == \"system\":\n",
    "            prompt_list.append(\"<|im_start|>assistant\\n\" + message[\"content\"] + \"\\n<|im_end|>\")\n",
    "        else:\n",
    "            raise ValueError(\"Invalid role: \" + message[\"role\"])\n",
    "    prompt_list.append(\"<|im_start|>user\\n-----请改写下面这个问题，让用户只关注该问题就能理解这个问题-----\\n\" + query + \"<|im_end|>\")\n",
    "    prompt = \"\\n\".join(prompt_list)\n",
    "    \n",
    "    try:\n",
    "        completion = client.chat.completions.create(model=qwen72b_modelId, messages=[{\"role\": \"user\", \"content\": prompt}])\n",
    "        rewritten_query = completion.choices[0].message.content\n",
    "        # from sentence_transformers import SentenceTransformer\n",
    "        sentences_old = query\n",
    "        sentences_new = rewritten_query\n",
    "        # model = SentenceTransformer('lier007/xiaobu-embedding-v2')\n",
    "        # embeddings_1 = model.encode(sentences_1, normalize_embeddings=True)\n",
    "        # embeddings_2 = model.encode(sentences_2, normalize_embeddings=True)\n",
    "        # similarity = embeddings_1 @ embeddings_2.T\n",
    "        # print(similarity)\n",
    "\n",
    "        context_embedding_response = client.embeddings.create(\n",
    "            model=embedding_model,\n",
    "            input=context,\n",
    "        )\n",
    "        context_embedding = context_embedding_response.data[0].embedding\n",
    "        query_embedding_response_1 = client.embeddings.create(\n",
    "            model=embedding_model,\n",
    "            input=sentences_old,\n",
    "        )\n",
    "        query_embedding_1 = query_embedding_response_1.data[0].embedding\n",
    "        query_embedding_response_2 = client.embeddings.create(\n",
    "            model=embedding_model,\n",
    "            input=sentences_new,\n",
    "        )\n",
    "        query_embedding_2 = query_embedding_response_2.data[0].embedding\n",
    "        similarity = query_embedding_1 @ query_embedding_2.T # 原问题和新问题的相似性\n",
    "        similarity_1 = query_embedding_1 @ context_embedding.T # 原问题和上下文的相似性\n",
    "        similarity_2 = query_embedding_2 @ context_embedding.T # 新问题和上下文的相似性\n",
    "\n",
    "        if similarity < 0.92 or similarity_2 > similarity_1: # 如果新问题与原问题相似度小于0.92或者新问题与上下文的相似度大于原问题与上下文的相似度，则改写\n",
    "            rewritten = True\n",
    "            rewrite_query = rewritten_query\n",
    "    except Exception as e:\n",
    "        print(f\"Error during OpenAI API call: {e}\")\n",
    "        return {\n",
    "            \"rewritten\": rewritten,\n",
    "            \"rewrite_query\": query\n",
    "        }\n",
    "    return {\n",
    "        \"rewritten\": rewritten,\n",
    "        \"rewrite_query\": rewrite_query\n",
    "    }\n",
    "\n",
    "\n",
    "# 此处不需要修改\n",
    "@app.route(\"/predict\",methods=[\"POST\"])\n",
    "def predict():\n",
    "    data = request.get_json()\n",
    "    messages = data.get('messages')\n",
    "\n",
    "    return digest_mrc(messages)\n",
    "\n",
    "\n",
    "# 此处不需要修改\n",
    "if __name__ == '__main__':\n",
    "   app.run(\"0.0.0.0\",80)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "# 第三次策略，判断是否需要重写查询。向量相似度判断+llm判断，改写类型有指代消岐和信息补全。主要在于如何判断是否需要改写。再考虑改写后的质量\n",
    "# 0.0.6\n",
    "from flask import Flask\n",
    "from flask import request\n",
    "from flask import jsonify\n",
    "# coding=UTF-8\n",
    "import json\n",
    "import os\n",
    "from flask import make_response\n",
    "from openai import OpenAI\n",
    " \n",
    " \n",
    "model_info = json.loads(os.getenv(\"LEADERBOARD_MODELHUB_KEY2INFO\")) # 环境变量名称固定\n",
    "token = model_info[\"token\"] # 访问modelhub模型的token，固定key\n",
    "base_url = model_info[\"entrypoint\"] # 访问modelhub模型的base url，固定key\n",
    "qwen72b_modelId = model_info[\"model_key2info\"][\"qwen72b\"][\"modelId\"] # key的qwen72b与提交配置里的自定义名称对应\n",
    "client = OpenAI(api_key=token, base_url=base_url)\n",
    "embedding_model = model_info[\"model_key2info\"][\"embedding\"][\"modelId\"]\n",
    "# res = client.chat.completions.create(model=qwen72b_modelId, **request.get_json())\n",
    "\n",
    "app = Flask(__name__)\n",
    "\n",
    "\n",
    "# 在这里实现你的策略\n",
    "def digest_mrc(messages):\n",
    "    \"\"\"\n",
    "    # Request:\n",
    "    {\n",
    "    \"messages\": [\n",
    "        {\"role\": \"user\", \"content\": \"今天上海的天气怎么样\"},\n",
    "        {\"role\": \"system\", \"content\": \"今天上海 18 度，晴天\"},\n",
    "        {\"role\": \"user\", \"content\": \"明天呢？\"},\n",
    "        {\"role\": \"system\", \"content\": \"明天预计 19 度，阴天\"},\n",
    "        {\"role\": \"user\", \"content\": \"广州什么天气？\"}\n",
    "    ]\n",
    "    }\n",
    "\n",
    "    # Response:\n",
    "    {\n",
    "    \"rewritten\": true/false\n",
    "    \"rewrite_query\": \"广州明天什么天气？\"\n",
    "    }\n",
    "    # 如果需要改写就是true，然后query就是改写后的问题，否则就是false，query字段可有可无\n",
    "    \"\"\"\n",
    "    rewritten = False\n",
    "    query = messages[-1][\"content\"]\n",
    "    context = \"\"\n",
    "    messages.pop(-1)\n",
    "    for message in messages:\n",
    "        context += message[\"content\"] + \"\\n\"\n",
    "    \n",
    "    try:\n",
    "        # from sentence_transformers import SentenceTransformer\n",
    "        sentences_old = query\n",
    "\n",
    "        context_embedding_response = client.embeddings.create(\n",
    "            model=embedding_model,\n",
    "            input=context,\n",
    "        )\n",
    "        context_embedding = context_embedding_response.data[0].embedding\n",
    "        query_embedding_response_1 = client.embeddings.create(\n",
    "            model=embedding_model,\n",
    "            input=sentences_old,\n",
    "        )\n",
    "        query_embedding_1 = query_embedding_response_1.data[0].embedding\n",
    "\n",
    "        similarity_1 = query_embedding_1 @ context_embedding.T # 原问题和上下文的相似性\n",
    "\n",
    "        if similarity_1 < 0.9:\n",
    "            rewritten = True\n",
    "            rewrite_query = query\n",
    "    except Exception as e:\n",
    "        print(f\"Error during OpenAI API call: {e}\")\n",
    "        return {\n",
    "            \"rewritten\": rewritten,\n",
    "            \"rewrite_query\": query\n",
    "        }\n",
    "    return {\n",
    "        \"rewritten\": rewritten,\n",
    "        \"rewrite_query\": rewrite_query\n",
    "    }\n",
    "\n",
    "\n",
    "# 此处不需要修改\n",
    "@app.route(\"/predict\",methods=[\"POST\"])\n",
    "def predict():\n",
    "    data = request.get_json()\n",
    "    messages = data.get('messages')\n",
    "\n",
    "    return digest_mrc(messages)\n",
    "\n",
    "\n",
    "# 此处不需要修改\n",
    "if __name__ == '__main__':\n",
    "   app.run(\"0.0.0.0\",80)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "# 0.0.7\n",
    "# 第三次策略，判断是否需要重写查询。向量相似度判断+llm判断，改写类型有指代消岐和信息补全。主要在于如何判断是否需要改写。再考虑改写后的质量\n",
    "from flask import Flask\n",
    "from flask import request\n",
    "from flask import jsonify\n",
    "# coding=UTF-8\n",
    "import json\n",
    "import os\n",
    "from flask import make_response\n",
    "from openai import OpenAI\n",
    " \n",
    " \n",
    "# res = client.chat.completions.create(model=qwen72b_modelId, **request.get_json())\n",
    "\n",
    "app = Flask(__name__)\n",
    "\n",
    "\n",
    "# 在这里实现你的策略\n",
    "def digest_mrc(messages):\n",
    "    \"\"\"\n",
    "    # Request:\n",
    "    {\n",
    "    \"messages\": [\n",
    "        {\"role\": \"user\", \"content\": \"今天上海的天气怎么样\"},\n",
    "        {\"role\": \"system\", \"content\": \"今天上海 18 度，晴天\"},\n",
    "        {\"role\": \"user\", \"content\": \"明天呢？\"},\n",
    "        {\"role\": \"system\", \"content\": \"明天预计 19 度，阴天\"},\n",
    "        {\"role\": \"user\", \"content\": \"广州什么天气？\"}\n",
    "    ]\n",
    "    }\n",
    "\n",
    "    # Response:\n",
    "    {\n",
    "    \"rewritten\": true/false\n",
    "    \"rewrite_query\": \"广州明天什么天气？\"\n",
    "    }\n",
    "    # 如果需要改写就是true，然后query就是改写后的问题，否则就是false，query字段可有可无\n",
    "    \"\"\"\n",
    "    rewritten = False\n",
    "    rewrite_query = messages[-1][\"content\"]\n",
    "    # 策略 1 直接语言模型 提示改写\n",
    "    ## 先将message处理成合适的prompt\n",
    "    query = messages[-1][\"content\"]\n",
    "    context = \"\"\n",
    "    messages.pop(-1)\n",
    "    for message in messages:\n",
    "        context += message[\"content\"] + \"\\n\"\n",
    "\n",
    "    \n",
    "    try:\n",
    "        # from sentence_transformers import SentenceTransformer\n",
    "        sentences_old = query\n",
    "        from sentence_transformers import SentenceTransformer\n",
    "        model = SentenceTransformer('moka-ai/m3e-base')\n",
    "        embeddings_old = model.encode(sentences_old, normalize_embeddings=True)\n",
    "        embeddings_context = model.encode(context, normalize_embeddings=True)\n",
    "        similarity_old_with_context = embeddings_old @ embeddings_context.T # 原问题和上下文的相似性\n",
    "\n",
    "        if similarity_old_with_context < 0.90:\n",
    "            rewritten = True\n",
    "            rewrite_query = query\n",
    "    except Exception as e:\n",
    "        print(f\"Error during OpenAI API call: {e}\")\n",
    "        return {\n",
    "            \"rewritten\": rewritten,\n",
    "            \"rewrite_query\": query\n",
    "        }\n",
    "    return {\n",
    "        \"rewritten\": rewritten,\n",
    "        \"rewrite_query\": rewrite_query\n",
    "    }\n",
    "\n",
    "\n",
    "# 此处不需要修改\n",
    "@app.route(\"/predict\",methods=[\"POST\"])\n",
    "def predict():\n",
    "    data = request.get_json()\n",
    "    messages = data.get('messages')\n",
    "\n",
    "    return digest_mrc(messages)\n",
    "\n",
    "\n",
    "# 此处不需要修改\n",
    "if __name__ == '__main__':\n",
    "   app.run(\"0.0.0.0\",80)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "# 第三次策略，判断是否需要重写查询。向量相似度判断+llm判断，改写类型有指代消岐和信息补全。主要在于如何判断是否需要改写。再考虑改写后的质量\n",
    "# 0.0.6\n",
    "from flask import Flask\n",
    "from flask import request\n",
    "from flask import jsonify\n",
    "# coding=UTF-8\n",
    "import json\n",
    "import os\n",
    "from flask import make_response\n",
    "from openai import OpenAI\n",
    " \n",
    " \n",
    "model_info = json.loads(os.getenv(\"LEADERBOARD_MODELHUB_KEY2INFO\")) # 环境变量名称固定\n",
    "token = model_info[\"token\"] # 访问modelhub模型的token，固定key\n",
    "base_url = model_info[\"entrypoint\"] # 访问modelhub模型的base url，固定key\n",
    "client = OpenAI(api_key=token, base_url=base_url)\n",
    "embedding_model = model_info[\"model_key2info\"][\"embedding\"][\"modelId\"]\n",
    "\n",
    "app = Flask(__name__)\n",
    "\n",
    "\n",
    "# 在这里实现你的策略\n",
    "def digest_mrc(messages):\n",
    "    \"\"\"\n",
    "    # Request:\n",
    "    {\n",
    "    \"messages\": [\n",
    "        {\"role\": \"user\", \"content\": \"今天上海的天气怎么样\"},\n",
    "        {\"role\": \"system\", \"content\": \"今天上海 18 度，晴天\"},\n",
    "        {\"role\": \"user\", \"content\": \"明天呢？\"},\n",
    "        {\"role\": \"system\", \"content\": \"明天预计 19 度，阴天\"},\n",
    "        {\"role\": \"user\", \"content\": \"广州什么天气？\"}\n",
    "    ]\n",
    "    }\n",
    "\n",
    "    # Response:\n",
    "    {\n",
    "    \"rewritten\": true/false\n",
    "    \"rewrite_query\": \"广州明天什么天气？\"\n",
    "    }\n",
    "    # 如果需要改写就是true，然后query就是改写后的问题，否则就是false，query字段可有可无\n",
    "    \"\"\"\n",
    "    rewritten = False\n",
    "    query = messages[-1][\"content\"]\n",
    "    context = \"\"\n",
    "    messages.pop(-1)\n",
    "    for message in messages:\n",
    "        context += message[\"content\"] + \"\\n\"\n",
    "    \n",
    "    try:\n",
    "        # from sentence_transformers import SentenceTransformer\n",
    "        sentences_old = query\n",
    "\n",
    "        context_embedding_response = client.embeddings.create(\n",
    "            model=embedding_model,\n",
    "            input=context,\n",
    "        )\n",
    "        context_embedding = context_embedding_response.data[0].embedding\n",
    "        query_embedding_response_1 = client.embeddings.create(\n",
    "            model=embedding_model,\n",
    "            input=sentences_old,\n",
    "        )\n",
    "        query_embedding_1 = query_embedding_response_1.data[0].embedding\n",
    "\n",
    "        similarity_1 = query_embedding_1 @ context_embedding.T # 原问题和上下文的相似性\n",
    "\n",
    "        if similarity_1 < 0.9:\n",
    "            rewritten = True\n",
    "            rewrite_query = query\n",
    "    except Exception as e:\n",
    "        print(f\"Error during OpenAI API call: {e}\")\n",
    "        return {\n",
    "            \"rewritten\": rewritten,\n",
    "            \"rewrite_query\": query\n",
    "        }\n",
    "    return {\n",
    "        \"rewritten\": rewritten,\n",
    "        \"rewrite_query\": rewrite_query\n",
    "    }\n",
    "\n",
    "\n",
    "# 此处不需要修改\n",
    "@app.route(\"/predict\",methods=[\"POST\"])\n",
    "def predict():\n",
    "    data = request.get_json()\n",
    "    messages = data.get('messages')\n",
    "\n",
    "    return digest_mrc(messages)\n",
    "\n",
    "\n",
    "# 此处不需要修改\n",
    "if __name__ == '__main__':\n",
    "   app.run(\"0.0.0.0\",80)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "相似度： 0.44943641652398214\n"
     ]
    }
   ],
   "source": [
    "import jieba\n",
    "from collections import Counter\n",
    "from sklearn.feature_extraction.text import TfidfVectorizer\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "\n",
    "def preprocess(text):\n",
    "    return ' '.join(jieba.cut(text))\n",
    "\n",
    "def calculate_cosine_similarity(question, context):\n",
    "    vectorizer = TfidfVectorizer()\n",
    "    tfidf_matrix = vectorizer.fit_transform([question] + [context])\n",
    "    cosine_similarities = cosine_similarity(tfidf_matrix[0:1], tfidf_matrix[1:])\n",
    "    return cosine_similarities[0][0]\n",
    "\n",
    "question = \"广州什么天气？\"\n",
    "context = \"广州\"\n",
    "\n",
    "preprocessed_question = preprocess(question)\n",
    "preprocessed_context = preprocess(context)\n",
    "\n",
    "similarity = calculate_cosine_similarity(preprocessed_question, preprocessed_context)\n",
    "print(\"相似度：\", similarity)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ChatTTS",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
